
from .dwpose import DWposeDetector
from .util import resize_image,HWC3
import cv2
from PIL import Image
import numpy as np


def get_pose(image:str|Image.Image|np.ndarray):
    dwprocessor = DWposeDetector()
    if isinstance(image,str):
        input_image = cv2.imread(image)
        input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
    elif isinstance(image,Image.Image):
        input_image = np.array(image)
    elif isinstance(image,np.ndarray):
        input_image = image
    else:
        raise ValueError("仅支持str路径或pil图像或np图像") 
    #print(input_image.shape)
    input_image = HWC3(input_image)
    img = resize_image(input_image, 384)
    #print(img.shape)
    #cv2.imwrite("../images/222-0.jpg",img)
    detected_map = dwprocessor(img)
    pil_pose = Image.fromarray(detected_map)
    #input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
    #cv2.imwrite("pose.png",detected_map)
    return pil_pose

